{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "\n",
    "import pybullet as p\n",
    "import pybullet_data\n",
    "from surrol.utils.pybullet_utils import (\n",
    "    step,\n",
    "    get_joints,\n",
    "    get_link_name,\n",
    "    reset_camera,\n",
    ")\n",
    "from surrol.robots.ecm import Ecm\n",
    "\n",
    "from surrol.const import ASSET_DIR_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling = 1.\n",
    "\n",
    "p.connect(p.GUI)\n",
    "# p.connect(p.DIRECT)\n",
    "p.setGravity(0, 0, -9.81)\n",
    "p.setAdditionalSearchPath(pybullet_data.getDataPath())\n",
    "reset_camera(yaw=10, pitch=-15, dist=0.9*scaling)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p.loadURDF(\"plane.urdf\", [0, 0, -0.001], globalScaling=1)\n",
    "\n",
    "# urdf_path = os.path.join(ASSET_DIR_PATH, 'ecm/ecm.urdf')\n",
    "# urdf_path = os.path.join(ASSET_DIR_PATH, \"ecm/ecm_dvrk.urdf\")\n",
    "\n",
    "# ecm = p.loadURDF(urdf_path, [0, 0, 0*scaling], \n",
    "#                  useFixedBase=True)\n",
    "\n",
    "ecm = Ecm([0, 0, 0*scaling], \n",
    "          scaling=scaling)\n",
    "ecm.reset_joint([0, 0, 0, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "joints = get_joints(ecm.body)\n",
    "print(\"There are {} joints.\\n\".format(len(joints)))\n",
    "\n",
    "for i in range(0, len(joints)):\n",
    "    print(get_link_name(ecm.body, i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# continously run\n",
    "p.setRealTimeSimulation(1)\n",
    "\n",
    "while True:\n",
    "    p.setGravity(0, 0, -9.81)\n",
    "    time.sleep(0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forward Kinematics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ecm.reset_joint([0, 0, 0, 0])\n",
    "joints = get_joints(ecm.body)\n",
    "\n",
    "pose_rcm = ecm.get_current_position()\n",
    "print(np.round(pose_rcm, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ecm.reset_joint([-0.0024, -0.0023, 0.0025, -0.0007])\n",
    "\n",
    "pose_rcm = ecm.get_current_position()\n",
    "print(np.round(pose_rcm, 4))\n",
    "\n",
    "# previously compute\n",
    "# [[ 1.0000  0.0006 -0.0024 -0.0000]\n",
    "#  [ 0.0007 -1.0000  0.0023  0.0000]\n",
    "#  [-0.0024 -0.0023 -1.0000 -0.0032]\n",
    "#  [ 0.0000  0.0000  0.0000  1.0000]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ecm.reset_joint([0.0884, -0.6098, 0.1961, -0.0118])\n",
    "\n",
    "pose_rcm = ecm.get_current_position()\n",
    "print(np.round(pose_rcm, 4))\n",
    "\n",
    "# [[ 0.9954  0.0623  0.0724  0.0142]\n",
    "#  [ 0.0097 -0.8197  0.5727  0.1127]\n",
    "#  [ 0.0950 -0.5694 -0.8166 -0.1607]\n",
    "#  [ 0.0000  0.0000  0.0000  1.0000]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ecm.reset_joint([0.5369, 0.1454, 0.0316, 0.1266])\n",
    "\n",
    "pose_rcm = ecm.get_current_position()\n",
    "print(np.round(pose_rcm, 4))\n",
    "\n",
    "# [[ 0.8431 -0.1820  0.5060  0.0163]\n",
    "#  [-0.1249 -0.9815  0.1449 -0.0047]\n",
    "#  [ 0.5231  0.0590 -0.8502 -0.0275]\n",
    "#  [ 0.0000  0.0000  0.0000  1.0000]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inverse Kinematics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_rcm = np.array([\n",
    "    [ 1.,  0.,  0.,  0.],\n",
    "    [ 0., -1.,  0.,  0.],\n",
    "    [ 0.,  0., -1., -0.0007],\n",
    "    [ 0.,  0.,  0.,  1.],\n",
    "])\n",
    "\n",
    "joints_inv = ecm.move(pose_rcm)\n",
    "step(0.5)\n",
    "print(np.round(joints_inv, 4))\n",
    "\n",
    "# [0, 0, 0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_rcm = np.array([\n",
    "    [ 1.0000,  0.0006, -0.0024, -0.0000],\n",
    "    [ 0.0007, -1.0000,  0.0023,  0.0000],\n",
    "    [-0.0024, -0.0023, -1.0000, -0.0032],\n",
    "    [ 0.0000,  0.0000,  0.0000,  1.0000],\n",
    "])\n",
    "\n",
    "joints_inv = ecm.move(pose_rcm)\n",
    "step(0.5)\n",
    "print(np.round(joints_inv, 4))\n",
    "\n",
    "# [-0.0024, -0.0023, 0.0025, -0.0007]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_rcm = np.array([\n",
    "    [ 0.9954,  0.0623,  0.0724,  0.0142],\n",
    "    [ 0.0097, -0.8197,  0.5727,  0.1127],\n",
    "    [ 0.0950, -0.5694, -0.8166, -0.1607],\n",
    "    [ 0.0000,  0.0000,  0.0000,  1.0000],\n",
    "])\n",
    "\n",
    "joints_inv = ecm.move(pose_rcm)\n",
    "step(0.5)\n",
    "print(np.round(joints_inv, 4))\n",
    "\n",
    "# [0.0884, -0.6098, 0.1961, -0.0118]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_rcm = np.array([\n",
    "    [ 0.8431, -0.182 ,  0.5061,  0.0163],\n",
    "    [-0.1249, -0.9815, -0.1449, -0.0047],\n",
    "    [ 0.5231,  0.0589, -0.8502, -0.0275],\n",
    "    [ 0.    ,  0.    ,  0.    ,  1.    ],\n",
    "])\n",
    "\n",
    "joints_inv = ecm.move(pose_rcm)\n",
    "step(0.5)\n",
    "print(np.round(joints_inv, 4))\n",
    "\n",
    "# [0.5369, 0.1454, 0.0316, 0.1266]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_rcm = np.array([\n",
    "    [ 0.84306447, -0.18201039,  0.50607759,  0.01634631],\n",
    "    [-0.12492979, -0.98152939, -0.14488822, -0.00467992],\n",
    "    [ 0.52310119,  0.05892595, -0.85023107, -0.02746246],\n",
    "    [ 0.        ,  0.        ,  0.        ,  1.        ],\n",
    "])\n",
    "\n",
    "joints_inv = ecm.move(pose_rcm)\n",
    "step(0.5)\n",
    "print(np.round(joints_inv, 4))\n",
    "\n",
    "# [0.5369, 0.1454, 0.0316, 0.1266]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "surrol",
   "language": "python",
   "name": "rorl"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
